// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

namespace kumo::pollux::functions::prestosql {

namespace {

class IPAddressCastTest : public functions::test::FunctionBaseTest {
 protected:
  std::optional<std::string> castToVarchar(
      const std::optional<std::string>& input) {
    auto result = evaluateOnce<std::string>(
        "cast(cast(c0 as ipaddress) as varchar)", input);
    return result;
  }

  std::optional<int128_t> castFromVarbinary(
      const std::optional<std::string>& input) {
    auto result =
        evaluateOnce<int128_t>("cast(from_hex(c0) as ipaddress)", input);
    return result;
  }

  std::optional<std::string> allCasts(const std::optional<std::string>& input) {
    auto result = evaluateOnce<std::string>(
        "cast(cast(cast(cast(c0 as ipaddress) as varbinary) as ipaddress) as varchar)",
        input);
    return result;
  }

  auto castToIPPrefixAndBackToIpVarchar(
      const std::optional<std::string>& input) {
    return evaluateOnce<std::string>(
        "cast(cast(cast(cast(cast(cast(c0 as ipaddress) as ipprefix) as varchar) as ipprefix) as ipaddress)  as varchar)",
        input);
  }
};

int128_t stringToInt128(const std::string& value) {
  int128_t res = 0;
  for (char c : value) {
    res = res * 10 + c - '0';
  }
  return res;
}

TEST_F(IPAddressCastTest, castToIPPrefix) {
  EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("1.2.3.4"), "1.2.3.4");
  EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("::ffff:1.2.3.4"), "1.2.3.4");
  EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("::ffff:102:304"), "1.2.3.4");
  EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("192.168.0.0"), "192.168.0.0");
  EXPECT_EQ(
      castToIPPrefixAndBackToIpVarchar(
          "2001:0db8:0000:0000:0000:ff00:0042:8329"),
      "2001:db8::ff00:42:8329");
  EXPECT_EQ(
      castToIPPrefixAndBackToIpVarchar("2001:db8:0:0:1:0:0:1"),
      "2001:db8::1:0:0:1");
  EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("::1"), "::1");
  EXPECT_EQ(
      castToIPPrefixAndBackToIpVarchar("2001:db8::ff00:42:8329"),
      "2001:db8::ff00:42:8329");
}

TEST_F(IPAddressCastTest, castToVarchar) {
  EXPECT_EQ(castToVarchar("::ffff:1.2.3.4"), "1.2.3.4");
  EXPECT_EQ(castToVarchar("0:0:0:0:0:0:13.1.68.3"), "::13.1.68.3");
  EXPECT_EQ(castToVarchar("1.2.3.4"), "1.2.3.4");
  EXPECT_EQ(castToVarchar("192.168.0.0"), "192.168.0.0");
  EXPECT_EQ(
      castToVarchar("2001:0db8:0000:0000:0000:ff00:0042:8329"),
      "2001:db8::ff00:42:8329");
  EXPECT_EQ(castToVarchar("2001:db8::ff00:42:8329"), "2001:db8::ff00:42:8329");
  EXPECT_EQ(castToVarchar("2001:db8:0:0:1:0:0:1"), "2001:db8::1:0:0:1");
  EXPECT_EQ(castToVarchar("2001:db8:0:0:1::1"), "2001:db8::1:0:0:1");
  EXPECT_EQ(castToVarchar("2001:db8::1:0:0:1"), "2001:db8::1:0:0:1");
  EXPECT_EQ(
      castToVarchar("2001:DB8::FF00:ABCD:12EF"), "2001:db8::ff00:abcd:12ef");
  POLLUX_ASSERT_THROW(
      castToVarchar("facebook.com"), "Invalid IP address 'facebook.com'");
  POLLUX_ASSERT_THROW(
      castToVarchar("localhost"), "Invalid IP address 'localhost'");
  POLLUX_ASSERT_THROW(
      castToVarchar("2001:db8::1::1"), "Invalid IP address '2001:db8::1::1'");
  POLLUX_ASSERT_THROW(
      castToVarchar("2001:zxy::1::1"), "Invalid IP address '2001:zxy::1::1'");
  POLLUX_ASSERT_THROW(
      castToVarchar("789.1.1.1"), "Invalid IP address '789.1.1.1'");
}

TEST_F(IPAddressCastTest, castFromVarbinary) {
  EXPECT_EQ(
      castFromVarbinary("00000000000000000000ffff01020304"),
      stringToInt128("281470698652420"));
  EXPECT_EQ(castFromVarbinary("01020304"), stringToInt128("281470698652420"));
  EXPECT_EQ(castFromVarbinary("c0a80000"), stringToInt128("281473913978880"));
  EXPECT_EQ(
      castFromVarbinary("20010db8000000000000ff0000428329"),
      stringToInt128("42540766411282592856904265327123268393"));
  EXPECT_THROW(castFromVarbinary("f000001100"), PolluxUserError);
}

TEST_F(IPAddressCastTest, allCasts) {
  EXPECT_EQ(allCasts("::ffff:1.2.3.4"), "1.2.3.4");
  EXPECT_EQ(
      allCasts("2001:0db8:0000:0000:0000:ff00:0042:8329"),
      "2001:db8::ff00:42:8329");
  EXPECT_EQ(allCasts("2001:db8::ff00:42:8329"), "2001:db8::ff00:42:8329");
}

TEST_F(IPAddressCastTest, nullTest) {
  EXPECT_EQ(castToVarchar(std::nullopt), std::nullopt);
  EXPECT_EQ(castFromVarbinary(std::nullopt), std::nullopt);
}

TEST_F(IPAddressCastTest, castRoundTrip) {
  auto strings = make_flat_vector<std::string>(
      {"87a0:ce14:8989:44c9:826e:b4d8:73f9:1542",
       "7cd6:bcec:1216:5c20:4b67:b1bd:173:ced",
       "192.128.0.0"});

  auto ipaddresses =
      evaluate("cast(c0 as ipaddress)", make_row_vector({strings}));
  auto stringsCopy =
      evaluate("cast(c0 as varchar)", make_row_vector({ipaddresses}));
  auto ipaddressesCopy =
      evaluate("cast(c0 as ipaddress)", make_row_vector({stringsCopy}));

  pollux::test::assertEqualVectors(strings, stringsCopy);
  pollux::test::assertEqualVectors(ipaddresses, ipaddressesCopy);
}
} // namespace

} // namespace kumo::pollux::functions::prestosql
